"""
基于 WebSocket 的语音处理服务
功能：
1. 中英文语音识别
2. 中英互译
使用技术：
- Flask-SocketIO 实现 WebSocket 通信
- FunASR 处理语音识别
- ModelScope 处理翻译任务
"""

import os
import logging
from typing import Any

from flask import Flask
from flask_socketio import SocketIO, emit
from funasr import AutoModel
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from flask_cors import CORS

# 初始化应用
app = Flask(__name__)
CORS(app)  # 允许跨域

# 配置 SocketIO
socketio = SocketIO(
    app,
    cors_allowed_origins="*",  # 允许所有来源连接
    async_mode='eventlet',  # 使用 eventlet 实现异步
    logger=True,  # 启用日志
    engineio_logger=True  # 显示 Engine.IO 日志
)

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# ==================== 模型初始化 ====================
# 注意：生产环境应该实现模型懒加载和单例管理

def init_models():
    """初始化所有需要的AI模型"""
    global model_zh, model_en, pipeline_trans_zh2en, pipeline_trans_en2zh

    logger.info("Loading Chinese ASR model...")
    model_zh = AutoModel(
        model="./used_weights/modelscope/hub/models/iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
        model_revision="v2.0.4",
        vad_model="./used_weights/modelscope/hub/models/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
        vad_model_revision="v2.0.4",
        punc_model="./used_weights/modelscope/hub/models/iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch",
        punc_model_revision="v2.0.4",
        disable_update=True,
        cache_dir="./used_weights",
        force_download=False,
        resume_download=False,
        # spk_model="cam++", spk_model_revision="v2.0.2",
        )

    # ======================= 语音转文字，英文语音 =======================
    model_dir = "./used_weights/modelscope/hub/iic/SenseVoiceSmall"

    model_en = AutoModel(
        model=model_dir,
        trust_remote_code=True,
        remote_code="./model.py",
        disable_update=True,
        vad_model="./used_weights/modelscope/hub/models/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
        vad_kwargs={"max_single_segment_time": 30000},
        device="cuda:0",
    )

    # ======================= 翻译，中文转英文 =======================
    pipeline_trans_zh2en = pipeline(task=Tasks.translation,
                                    model="./used_weights/modelscope/hub/models/damo/nlp_csanmt_translation_zh2en")

    # ======================= 翻译，英文转中文 =======================
    pipeline_trans_en2zh = pipeline(task=Tasks.translation,
                                    model="./used_weights/modelscope/hub/models/damo/nlp_csanmt_translation_en2zh")


# ==================== WebSocket 事件处理 ====================
from flask import Flask, request as flask_request

# 类型声明
socketio_request: Any = flask_request


@socketio.on('connect')
def handle_connect():
    """客户端连接事件"""
    client_id = socketio_request.sid
    logger.info(f'New client connected: {client_id}')
    emit('connection_response', {
        'status': 'connected',
        'client_id': client_id
    })


@socketio.on('disconnect')
def handle_disconnect():
    """客户端断开事件"""
    logger.info(f'Client disconnected: {socketio_request.sid}')


@socketio.on('speech_to_text')
def handle_speech_to_text(data):
    """
    处理语音转文本请求
    参数格式:
    {
        "language": "zh|en",  # 语言类型
        "chunk": binary_data,  # 音频二进制数据
        "is_final": bool,      # 是否最后一块
        "file_id": str         # 文件唯一ID
    }
    """
    temp_file = None
    try:
        file_id = data.get('file_id')
        language = data.get('language')
        chunk = data.get('chunk')
        is_final = data.get('is_final', False)

        # 验证参数
        if not all([file_id, language, chunk]):
            raise ValueError("Missing required parameters")

        # 保存分片到临时文件
        temp_file = os.path.join('uploads', f"{file_id}.wav")
        with open(temp_file, 'ab') as f:
            f.write(chunk)

        if is_final:
            logger.info(f"Processing speech recognition for {file_id}")

            # 根据语言选择模型
            if language == 'zh':
                result = model_zh.generate(input=temp_file)
            elif language == 'en':
                result = model_en.generate(input=temp_file)
            else:
                raise ValueError("Unsupported language")

            # 发送结果
            emit('speech_result', {
                'status': 'success',
                'text': result[0]['text'],
                'file_id': file_id
            })

            # 清理临时文件
            os.remove(temp_file)
        else:
            # 发送进度更新
            emit('progress_update', {
                'file_id': file_id,
                'progress': os.path.getsize(temp_file)
            })

    except Exception as e:
        logger.error(f"Speech recognition error: {str(e)}")
        emit('error', {
            'error': str(e),
            'file_id': data.get('file_id')
        })
        if temp_file and os.path.exists(temp_file):
            os.remove(temp_file)


@socketio.on('translate_text')
def handle_translation(data):
    """
    处理文本翻译请求
    参数格式:
    {
        "text": "要翻译的文本",
        "source_lang": "zh|en",
        "target_lang": "zh|en"
    }
    """
    try:
        text = data.get('text')
        source_lang = data.get('source_lang')
        target_lang = data.get('target_lang')

        # 验证参数
        if not all([text, source_lang, target_lang]):
            raise ValueError("Missing required parameters")

        # 选择正确的翻译管道
        if source_lang == 'zh' and target_lang == 'en':
            result = pipeline_trans_zh2en(input=text)
        elif source_lang == 'en' and target_lang == 'zh':
            result = pipeline_trans_en2zh(input=text)
        else:
            raise ValueError("Unsupported language pair")

        # 发送翻译结果
        emit('translation_result', {
            'status': 'success',
            'original_text': text,
            'translated_text': result['translation'],
            'language_pair': f"{source_lang}-{target_lang}"
        })

    except Exception as e:
        logger.error(f"Translation error: {str(e)}")
        emit('error', {'error': str(e)})


# ==================== 启动服务 ====================

if __name__ == '__main__':
    # 创建上传目录
    os.makedirs('uploads', exist_ok=True)

    # 初始化模型
    init_models()

    # 启动服务
    logger.info("Starting WebSocket server...")
    socketio.run(
        app,
        host='0.0.0.0',
        port=5072,
        debug=False,
        use_reloader=False
    )
